import random
import torch
from torch.autograd import Variable

"""
PyTorch: 控制流 + 共享权重

作为动态图和权值共享的例子，我们实现了一个不同的模型： 一个全连接的 ReLU 网络，每次前向传播时， 从1到4随机选择一个数来作为隐藏层的层数。=，多次重复使用相同的权重来计算最内层的隐藏层。
对于这个模型，我们可以使用普通的 Python 控制流来实现循环，在定义前向传播时，通过简单的重复使用同一个模块，我们可以实现最内层之间的权重共享。
"""

# 动态网络
class DynamicNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        在构造器中我们构造三个nn.Linear实例用于正向传播
        In the constructor we construct three nn.Linear instances that we will use
        in the forward pass.
        """
        super(DynamicNet, self).__init__()
        self.input_linear = torch.nn.Linear(D_in, H)
        self.middle_linear = torch.nn.Linear(H, H)
        self.output_linear = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        对于模型的正向传播，我们随机的选择任何一个0，1，2，3次重复使用middle_linear模块多次计算隐藏层
        For the forward pass of the model, we randomly choose either 0, 1, 2, or 3
        and reuse the middle_linear Module that many times to compute hidden layer
        representations.

        因为每一个正向传播都创建一个动态计算图，因此当以模型正向传播时我们可以使用正常的像
        循环或条件语句类的控制流操作。
        Since each forward pass builds a dynamic computation graph, we can use normal
        Python control-flow operators like loops or conditional statements when
        defining the forward pass of the model.

        这里我们也看到当定义一个计算图时可以完美安全的多次使用相同的模块
        这对于每个模型都只能用一次的Lua Torch是一个很大的改善
        Here we also see that it is perfectly safe to reuse the same Module many
        times when defining a computational graph. This is a big improvement from Lua
        Torch, where each Module could be used only once.
        """

        h_relu = self.input_linear(x).clamp(min=0)
        for _ in range(random.randint(0, 3)):
            h_relu = self.middle_linear(h_relu).clamp(min=0)
        y_pred = self.output_linear(h_relu)
        return y_pred

# N 批的大小
# D_in 输入的数据纬度
# H 隐藏层纬度
# D_out 输出数据纬度
# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# 创建随机的张量保存输入和输出，放到变量里
# Create random Tensors to hold inputs and outputs, and wrap them in Variables
x = Variable(torch.randn(N, D_in))
y = Variable(torch.randn(N, D_out), requires_grad=False)

# 用实例化上面的类构造我们的模型
# Construct our model by instantiating the class defined above
model = DynamicNet(D_in, H, D_out)

# 构造我们的损失函数和优化器。使用朴素SGD训练这个奇怪的模型时很困难的，因此我们使用冲量
# Construct our loss function and an Optimizer. Training this strange model with
# vanilla stochastic gradient descent is tough, so we use momentum
criterion = torch.nn.MSELoss(size_average=False)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
for t in range(500):
    # 正向传播： 把x传递到模型计算预测的y值
    # Forward pass: Compute predicted y by passing x to the model
    y_pred = model(x)

    # 计算和打印损失
    # Compute and print loss
    loss = criterion(y_pred, y)
    print(t, loss.data[0])

    # 清零梯度，执行后向传播，更新权重
    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()